import cv2
from insightface.app import FaceAnalysis
import torch
from diffusers import StableDiffusionXLPipeline, DDIMScheduler, AutoencoderKL
from PIL import Image
from ip_adapter.ip_adapter_faceid import IPAdapterFaceIDXL


app = FaceAnalysis(name="buffalo_l", providers=['CUDAExecutionProvider', 'CPUExecutionProvider'])
app.prepare(ctx_id=0, det_size=(320, 320))

image = cv2.imread("tmp/person1.png")
faces = app.get(image)
faceid_embeds = torch.from_numpy(faces[0].normed_embedding).unsqueeze(0)


repo_id = "SG161222/RealVisXL_V3.0"
ip_ckpt = "/home/aigc/.cache/huggingface/hub/models--h94--IP-Adapter-FaceID/snapshots/c36a25fd18f826b122f1c3162a7af961d5861aa3/ip-adapter-faceid_sdxl.bin"
scheduler = DDIMScheduler(
    num_train_timesteps=1000,
    beta_start=0.00085,
    beta_end=0.012,
    beta_schedule="scaled_linear",
    clip_sample=False,
    set_alpha_to_one=False,
    steps_offset=1,
)

pipe = StableDiffusionXLPipeline.from_pretrained(
    repo_id,
    torch_dtype=torch.float16,
    scheduler=scheduler,
    add_watermarker=False
)

# load ip-adapter
ip_model = IPAdapterFaceIDXL(pipe, ip_ckpt, device="cuda")

# generate image
prompt = "A closeup shot of a beautiful Asian teenage girl in a white dress wearing small silver earrings in the garden, under the soft morning light"
negative_prompt = "monochrome, lowres, bad anatomy, worst quality, low quality, blurry"

output_image = ip_model.generate(
    prompt=prompt, 
    negative_prompt=negative_prompt,
    faceid_embeds=faceid_embeds,
    num_samples=4, 
    num_inference_steps=30, 
    scale=1, 
    #seed=2023,
    width=1024, 
    height=1024
)[0]

output_image.save("tmp/c05.png")

